import os
import sys

from openai import OpenAI
sys.path.append('/JiuTian/swarm/swarm-main/swarm')
from swarm import Swarm,Agent

def process_and_print_streaming_response(response):
    content = ""
    last_sender = ""

    # 处理响应中的每一个片段
    for chunk in response:
        if "sender" in chunk:
            last_sender = chunk["sender"]  # 保存消息发送者的名字

        if "content" in chunk and chunk["content"] is not None:
            # 如果当前内容为空并且有消息发送者，输出发送者名字
            if not content and last_sender:
                print(f"\033[94m{last_sender}:\033[0m", end=" ", flush=True)
                last_sender = ""
            # 输出消息内容
            print(chunk["content"], end="", flush=True)
            content += chunk["content"]

        if "tool_calls" in chunk and chunk["tool_calls"] is not None:
            # 处理工具调用
            for tool_call in chunk["tool_calls"]:
                f = tool_call["function"]
                name = f["name"]
                if not name:
                    continue
                # 输出工具调用的函数名
                print(f"\033[94m{last_sender}: \033[95m{name}\033[0m()")

        if "delim" in chunk and chunk["delim"] == "end" and content:
            # 处理消息结束的情况，换行表示结束
            print()  # End of response message
            content = ""

        if "response" in chunk:
            # 返回最终的完整响应
            return chunk["response"]
#测试基本功能
client = OpenAI(api_key=os.environ["DEEPSEEK_API_KEY"], base_url=os.environ["DEEPSEEK_BASE_URL"])

def transfer_to_agent_b():
    return agent_b


model_name = "deepseek-chat"
#测试多智能体功能
swarm_client = Swarm(client)
agent_a = Agent(name = "Agent A",instructions= "你是一个乐于助人的助手",functions=[transfer_to_agent_b],model=model_name)
agent_b = Agent(name = "Agent B",instructions= "只使用俳句回答。",model=model_name)

response = swarm_client.run(agent=agent_a,messages=[{"role":"user","content":"我想和智能体B谈话，请问你你喜欢红楼梦吗？"}],stream=True)

responseText = process_and_print_streaming_response(response)
print(responseText)